import argparse
import os
import numpy as np
import matplotlib.pyplot as plt

# 定义区间和标签
bins = [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
labels = ['[0,0]', '(0,10]', '(10,20]', '(20,30]', '(30,40]', '(40,50]', 
          '(50,60]', '(60,70]', '(70,80]', '(80,90]', '(90,100)', '[100,100]']

# 定义处理TSV文件的函数，不使用csv模块
def read_tsv(file_path):
    data = []
    with open(file_path, 'r') as file:
        # 读取第一行，忽略其余内容
        file.readline()
        for line in file:
            # 将每一行用制表符分割，并将第一列转换为 float，后两列转换为 int
            row = line.strip().split('\t')
            data.append([float(row[0]), float(row[4]), float(row[5])])
    return data

# 按照区间划分数据
def classify_and_count(data):
    bin_data = {label: {'col2': [0]*101, 'col3': [0]*101} for label in labels}  # 初始化0-100频次的计数
    for row in data:
        value, col2, col3 = row
        value = int(value*100)
        col2 = int(col2*100)
        col3 = int(col3*100)

        # 单独处理 [0,0] 和 [100,100] 区间的边界情况
        if value == 0:
            if 0 <= col2 <= 100:
                bin_data['[0,0]']['col2'][col2] += 1
            if 0 <= col3 <= 100:
                bin_data['[0,0]']['col3'][col3] += 1
        elif value == 100:
            if 0 <= col2 <= 100:
                bin_data['[100,100]']['col2'][col2] += 1
            if 0 <= col3 <= 100:
                bin_data['[100,100]']['col3'][col3] += 1
        else:
            # 处理其他区间的情况
            for i in range(1, len(bins) - 1):
                if bins[i] < value <= bins[i + 1]:
                    if 0 <= col2 <= 100:
                        bin_data[labels[i]]['col2'][col2] += 1
                    if 0 <= col3 <= 100:
                        bin_data[labels[i]]['col3'][col3] += 1
                    break
    return bin_data

# 绘制并保存柱状图
def plot_histogram(bin_label, col2_freq, col3_freq, output_dir):
    x = np.arange(101)  # x轴上的0到100
    width = 0.35  # 柱子的宽度

    fig, ax = plt.subplots()

    # 绘制col2的柱状图，偏移0
    ax.bar(x - width/2, col2_freq, width, label='Rockfish', color='red')

    # 绘制col3的柱状图，偏移width
    ax.bar(x + width/2, col3_freq, width, label='T5', color='green')

    # 设置图的标题和标签
    ax.set_xlabel('Value')
    ax.set_ylabel('Frequency')
    ax.set_title(f'Frequency in Range {bin_label}')
    ax.legend()

    # 保存图片
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, f'{bin_label}.png'))
    plt.close()


# 主函数
def main():
    # 使用argparse解析输入参数
    parser = argparse.ArgumentParser(description="Process TSV file and generate histograms.")
    parser.add_argument("input_file", type=str, help="Path to the input TSV file")
    parser.add_argument("output_dir", type=str, help="Directory to save the output images")
    args = parser.parse_args()

    # 检查输出目录是否存在，不存在则创建
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    # 读取TSV文件数据
    data = read_tsv(args.input_file)

    # 将数据分类到不同区间
    bin_data =classify_and_count(data)

    # 为每个区间生成单独的柱状图并保存
    for bin_label, freqs in bin_data.items():
        plot_histogram(bin_label, freqs['col2'], freqs['col3'], args.output_dir)

if __name__ == "__main__":
    main()
